import pandas as pd
import torch
from torch import nn
from torch.utils.data import DataLoader
import torch.optim as optim
from dataset import NewsRatingDataset
from model import RecModel
import pickle as pkl
from tqdm import tqdm
import os
from recInterface import saveNewsAndUserFeature
from recInterface import getKNNitem, getUserMostLike


def train(model, dataset, n_epochs):
    datasets = NewsRatingDataset(df=dataset)
    dataloader = DataLoader(datasets, batch_size=64, shuffle=True)

    device = torch.device("cuda:0")

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    learning_rate = 1e-3

    # 损失函数, 均方误差
    loss_function = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    for epoch in range(n_epochs):
        # train
        loss_all = 0
        pbar = tqdm(total=len(dataloader))
        for _, sample_batch in enumerate(dataloader):
            user_input = sample_batch['user_input']
            news_input = sample_batch['news_input']
            target = sample_batch['target'].to(device)

            model.zero_grad()

            tag_rank, _, _ = model(user_input, news_input)
            loss = loss_function(tag_rank, target)
            # if i_batch % 20 == 0:
            #     print('loss: {}'.format(loss), end='\r')
            loss_all += loss
            loss.backward()
            optimizer.step()

            pbar.update()
        pbar.close()
        tqdm.write("Epoch {}:\t aver_loss: {}".format(epoch, loss_all / len(dataloader)))


if __name__ == '__main__':
    if not os.path.exists('Params'):
        os.mkdir('Params')

    dataset = pd.read_pickle('./data/data.pkl')
    n_user = pkl.load(open('./data/n_user.pkl', 'rb'))
    n_news = pkl.load(open('./data/n_news.pkl', 'rb'))
    news_title_dict = pkl.load(open('./data/news_title_dict.pkl', 'rb'))

    device = torch.device("cuda:0")
    model = RecModel(
        n_user=n_user + 1,
        n_news=n_news + 1,
        n_age=7,
        n_job=21,
        n_news_type=15 + 1,
        n_news_title_dict=len(news_title_dict) + 1,
        n_news_title_maxlen=48,
        cnn_kernel_size_list={2, 3, 4, 5},
        cnn_kernel_num=8,
        device=device).to(device)

    # train model
    train(model, dataset, 1)
    torch.save(model.state_dict(), 'Params/model_params.pkl')

    # get user and news feature
    # model.load_state_dict(torch.load('Params/model_params.pkl'))
    # saveNewsAndUserFeature(model=model)

    # test recsys
    print(getKNNitem(itemID=100,itemName='news',K=10))  #推荐相似用户或者相似新闻
    print(getUserMostLike(uid=100)) #推荐用户新闻
